from ultralytics import YOLO
import serial
import torch
from collections import deque,defaultdict

class RCDecisionMaker:
  def __init__(self,OpenSerial=False,bufferSize=10):
    if OpenSerial is True: 
      try: 
          self.serial = serial.Serial(port="/dev/ttyUSB0",baudrate=115200)
          print("Success in init serial.")
      except:
          self.serial = serial.Serial(port="/dev/ttyUSB1",baudrate=115200)
          print("Success in init serial.")

    else:
       self.serial = None

    self.bufferSize = bufferSize
    self.frameCounter = 0
    
    # 每个不同的 class_name都会自动对应一个自己的固定长度为 bufferSize 的队列（deque）。
    self.class_buffers = defaultdict((lambda: deque(maxlen=self.bufferSize))) 
    
    # 防止一个类别连续出现时反复发送同一个指令，直到即 buffer 不再全部是这个类，再重新进入判定逻辑。
    self.last_sent_classes = set() 
    self.last_seen_frames = defaultdict(int)  # 记录类别最后出现的帧号

  def MakeDecision(self,YOLOResultBox,names,frameCounter,zDepth):
    """
      names 是所有类别
    """
    if self.serial is None:
      print("serial does not exist.")
    self.frameCounter = frameCounter # need to modify
    #print("Judging YOLOBox...")
    # print("YOLOResultBox inside MakeDecision:", YOLOResultBox)
    if YOLOResultBox.data.shape[0] == 0:
      current_frame = frameCounter
      print(f'current_frame:',current_frame)
      for class_name in list(self.last_seen_frames.keys()):
          if current_frame - self.last_seen_frames[class_name] >= 10:
              self.class_buffers[class_name].clear()
              if class_name in self.last_sent_classes:
                  self.last_sent_classes.remove(class_name)
              del self.last_seen_frames[class_name]
      return

    class_id = int(YOLOResultBox.cls[0].cpu().item())
    #print("class_id:", class_id, type(class_id))  # 检查 class_id 类型 
    class_name = names[class_id]
    self.last_seen_frames[class_name] = self.frameCounter
    self.class_buffers[class_name].append(class_name)
    buffer = self.class_buffers[class_name]


    if len(buffer) == self.bufferSize and all(item == class_name for item in buffer):
      # print('zero')
      # print(f"Class: {class_name}, Buffer: {list(self.class_buffers[class_name])}, Last sent: {self.last_sent_classes}")
      if class_name not in self.last_sent_classes:
        if self._send_message(class_name,zDepth):
          
          self.last_sent_classes.add(class_name)
          print("NOT MY Fault! I have sent message.%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
        print('good but not send.')
    else:
      if class_name in self.last_sent_classes:
        self.last_sent_classes.remove(class_name)
    
    #print('fine') 
    # 第三阶段：清理30帧未出现的类别
    current_frame = self.frameCounter
    for class_name in list(self.last_seen_frames.keys()):
        if current_frame - self.last_seen_frames[class_name] >= 30:
            self.class_buffers[class_name].clear()
            if class_name in self.last_sent_classes:
                self.last_sent_classes.remove(class_name)
            del self.last_seen_frames[class_name]

  def _send_message(self,class_name,zDepth):
    print('================================send message=================================')
    if self.serial:

      zDepth_str = f",{zDepth:.2f}\n"  
      try:
          if class_name == "redball" and zDepth<=0.93:
            print(f'&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&')
            self.serial.write('R'.encode('utf-8'))
            print(f"Sent to serial: {'R'.strip()}")
            return True

          elif class_name == "blueball" and zDepth<=0.60:
            self.serial.write('B'.encode('utf-8'))
            print(f"Sent to serial: {'B'.strip()}")
            return True

          elif class_name == "purpleball" and zDepth<=0.30:
            self.serial.write('P'.encode('utf-8'))
            print(f"Sent to serial: {'P'.strip()}")
            return True 

          elif class_name == "one" and zDepth<=0.60:
            self.serial.write('1'.encode('utf-8'))
            print(f"Sent to serial: {'1'.strip()}")
            return True

          elif class_name == "two" and zDepth<=0.60:
            self.serial.write('2'.encode('utf-8'))
            print(f"Sent to serial: {'2'.strip()}")
            return True
          
          elif class_name == "three" and zDepth<=0.60:
            self.serial.write('3'.encode('utf-8'))
            print(f"Sent to serial: {'3'.strip()}")
            return True

          elif class_name == "four" and zDepth<=0.60:
            self.serial.write('4'.encode('utf-8'))
            print(f"Sent to serial: {'4'.strip()}")
            return True

          elif class_name == "five" and zDepth<=0.60:
            self.serial.write('5'.encode('utf-8'))
            print(f"Sent to serial: {'5'.strip()}")
            return True

          elif class_name == "six" and zDepth<=0.60:
            self.serial.write('6'.encode('utf-8'))
            print(f"Sent to serial: {'6'.strip()}")
            return True
          
          else:
              print('zDepths not sactisfy.')
              return False
#         self.serial.write(zDepth_str.encode('utf-8'))
#         print(f"Sent zDepth to serial: {zDepth_str.strip()}")
      except serial.SerialException as e:
                print(f"Serial write error: {e}")
    else:
      print(f"DEBUG MODE - Detected: {class_name}, zDepth: {zDepth}")
      return True

